/**
 * @file main.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include <cstdint>
#include <iostream>

#include "acl/acl.h"
#include "common.h"
#include "op_runner.h"

bool g_isDevice = false;
int deviceId = 0;

OperatorDesc CreateOpDesc()
{
    int case_id = 9;

    // 声明变量
    std::vector<int64_t> shape;
    std::vector<int64_t> maskShape;
    aclDataType dataType;

    // 固定项
    std::vector<int64_t> elemShape{1};
    aclDataType maskType = ACL_BOOL;
    aclFormat format = ACL_FORMAT_ND;

    switch (case_id) {
        case 0:
            shape = {512, 512, 512};
            maskShape = {1, 1, 512};
            dataType = ACL_FLOAT16;
            break;
        case 1:
            shape = {4096, 2048};
            maskShape = {4096, 1};
            dataType = ACL_FLOAT;
            break;
        case 2:
            shape = {1, 1, 1, 1, 100000};
            maskShape = {1, 1, 1, 1, 1};
            dataType = ACL_INT32;
            break;
        case 3:
            shape = {65536};
            maskShape = {1}; // 0D标量掩码，与 Python 空列表对应
            dataType = ACL_UINT8;
            break;
        case 4:
            shape = {1};
            maskShape = {1};
            dataType = ACL_BOOL;
            break;
        case 5:
            shape = {16, 1, 1, 1, 1, 65535};
            maskShape = {1, 1, 1, 1, 1, 65535};
            dataType = ACL_INT64;
            break;
        case 6:
            shape = {2048, 2048, 2};
            maskShape = {2048, 1, 1};
            dataType = ACL_FLOAT; // 用float32模拟bfloat16，与 Python 一致
            break;
        case 7:
            shape = {1, 1638, 1, 1638};
            maskShape = {1, 1, 1, 1638};
            dataType = ACL_DOUBLE;
            break;
        case 8:
            shape = {32, 1, 8192};
            maskShape = {32, 1, 1};
            dataType = ACL_INT16;
            break;
        case 9:
            shape = {1, 1, 1, 1, 1, 1, 262144};
            maskShape = {1, 1, 1, 1, 1, 1, 1};
            dataType = ACL_UINT64;
            break;
        default:
            throw std::invalid_argument("不支持的case_id，仅支持0-9");
            break;
    }
    OperatorDesc opDesc;
    // 有几个输入就AddInputTensorDesc几次
    opDesc.AddInputTensorDesc(dataType, shape.size(), shape.data(), format);
    opDesc.AddInputTensorDesc(maskType, maskShape.size(), maskShape.data(), format);
    opDesc.AddInputTensorDesc(dataType, elemShape.size(), elemShape.data(), format);
    opDesc.AddOutputTensorDesc(dataType, shape.size(), shape.data(), format);
    return opDesc;
}

bool SetInputData(OpRunner &runner)
{
    size_t fileSize = 0;
    // 确认输入个数
    ReadFile("../input/input_x.bin", fileSize, runner.GetInputBuffer<void>(0), runner.GetInputSize(0));
    ReadFile("../input/mask.bin", fileSize, runner.GetInputBuffer<void>(1), runner.GetInputSize(1));
    ReadFile("../input/fill_value.bin", fileSize, runner.GetInputBuffer<void>(2), runner.GetInputSize(2));
    INFO_LOG("Set input success");
    return true;
}

bool ProcessOutputData(OpRunner &runner)
{
    WriteFile("../output/output_z.bin", runner.GetOutputBuffer<void>(0), runner.GetOutputSize(0));
    INFO_LOG("Write output success");
    return true;
}

void DestroyResource()
{
    bool flag = false;
    if (aclrtResetDevice(deviceId) != ACL_SUCCESS) {
        ERROR_LOG("Reset device %d failed", deviceId);
        flag = true;
    }
    INFO_LOG("Reset Device success");
    if (aclFinalize() != ACL_SUCCESS) {
        ERROR_LOG("Finalize acl failed");
        flag = true;
    }
    if (flag) {
        ERROR_LOG("Destroy resource failed");
    } else {
        INFO_LOG("Destroy resource success");
    }
}

bool InitResource()
{
    std::string output = "../output";
    if (access(output.c_str(), 0) == -1) {
        int ret = mkdir(output.c_str(), 0700);
        if (ret == 0) {
            INFO_LOG("Make output directory successfully");
        } else {
            ERROR_LOG("Make output directory fail");
            return false;
        }
    }

    // acl.json is dump or profiling config file
    if (aclInit("../scripts/acl.json") != ACL_SUCCESS) {
        ERROR_LOG("acl init failed");
        return false;
    }

    if (aclrtSetDevice(deviceId) != ACL_SUCCESS) {
        ERROR_LOG("Set device failed. deviceId is %d", deviceId);
        (void)aclFinalize();
        return false;
    }
    INFO_LOG("Set device[%d] success", deviceId);

    // runMode is ACL_HOST which represents app is running in host
    // runMode is ACL_DEVICE which represents app is running in device
    aclrtRunMode runMode;
    if (aclrtGetRunMode(&runMode) != ACL_SUCCESS) {
        ERROR_LOG("Get run mode failed");
        DestroyResource();
        return false;
    }
    g_isDevice = (runMode == ACL_DEVICE);
    INFO_LOG("Get RunMode[%d] success", runMode);

    return true;
}

bool RunOp()
{
    // create op desc
    OperatorDesc opDesc = CreateOpDesc();

    // create Runner
    OpRunner opRunner(&opDesc);
    if (!opRunner.Init()) {
        ERROR_LOG("Init OpRunner failed");
        return false;
    }

    // Load inputs
    if (!SetInputData(opRunner)) {
        ERROR_LOG("Set input data failed");
        return false;
    }

    // Run op
    if (!opRunner.RunOp()) {
        ERROR_LOG("Run op failed");
        return false;
    }

    // process output data
    if (!ProcessOutputData(opRunner)) {
        ERROR_LOG("Process output data failed");
        return false;
    }

    INFO_LOG("Run op success");
    return true;
}

int main(int argc, char **argv)
{
    if (!InitResource()) {
        ERROR_LOG("Init resource failed");
        return FAILED;
    }
    INFO_LOG("Init resource success");

    if (!RunOp()) {
        DestroyResource();
        return FAILED;
    }

    DestroyResource();

    return SUCCESS;
}
